# -*- coding: utf-8 -*-
"""
Created on Mon Nov 21 12:54:55 2022

@author: sthgy
"""


import datetime
import numpy as np
import pandas as pd
import akshare as ak
from functools import lru_cache


@lru_cache()
def get_fund_list(fund_type=None):
    ''' 获取指定类型的基金，不指定则获取全部基金 '''
    with pd.HDFStore('./FundData/data.h5') as hdf_file:
        fund_info = pd.read_hdf(hdf_file, 'fund_info')
    if not fund_type:
        # 剔除了货币、商品、REITs、FOF等类型 
        fund_sift = fund_info.loc[fund_info['基金类型'].apply(lambda x: x in [
            '混合型-灵活',
            '债券型-可转债',
            '债券型-长债',
            '混合型-偏股',
            '指数型-股票',
            '债券型-混合债',
            '股票型',
            '债券型-中短债',
            '混合型-偏债',
            '混合-绝对收益',
            '混合型-平衡'
            ]), :]
        return fund_sift
    elif fund_type in ['债券型', '股票型']:
        # 基金大类
        fund_sift = fund_info.loc[fund_info['基金类型'].apply(
            lambda x: fund_type in x), :]
        return fund_sift
    elif fund_type in ['混合', '混合型']:
        # 基金大类
        fund_sift = fund_info.loc[fund_info['基金类型'].apply(
            lambda x: '混合' in x), :]
        return fund_sift

def date_format(date):
    ''' 调整时间格式 '''
    if isinstance(date, datetime.date):
        return date
    elif isinstance(date, datetime.datetime):
        return date.date()
    elif isinstance(date, int) or isinstance(date, float):
        date_string = str(int(date))
        return datetime.datetime.strptime(date_string, '%Y%m%d').date()
    elif isinstance(date, str):
        if '-' in date:
            return datetime.datetime.strptime(date, '%Y-%m-%d').date()
        elif str(date).isnumeric():
            return datetime.datetime.strptime(date, '%Y%m%d').date()
    else:
        raise ValueError(f'不支持的日期格式：{date}')

@lru_cache()
def get_trade_dates_all():
    ''' 获取全部交易日（补一个缓存） '''
    return ak.tool_trade_date_hist_sina()['trade_date']

def get_trade_dates(start_date, end_date):
    ''' 获取日期序列 '''
    trade_dates = get_trade_dates_all()
    dates_sifted = trade_dates[
        (trade_dates >= date_format(start_date))&
        (trade_dates <= date_format(end_date))
        ].sort_values().reset_index(drop=True)
    return dates_sifted

@lru_cache()
def read_info_all():
    ''' 从本地文件读取全部基金信息 '''
    with pd.HDFStore('./FundData/data.h5', 'r') as hdf_file:
        info_all = pd.read_hdf(hdf_file, 'fund_info')
    return info_all

@lru_cache()
def read_nav_all():
    ''' 从本地文件读取全部（累计）净值序列 '''
    with pd.HDFStore('./FundData/data.h5', 'r') as hdf_file:
        nav_all = pd.read_hdf(hdf_file, 'fund_cumnav')
    return nav_all

@lru_cache()
def read_scale_all():
    ''' 从本地文件读取全部基金规模序列 '''
    with pd.HDFStore('./FundData/data.h5', 'r') as hdf_file:
        fund_scale = pd.read_hdf(hdf_file, 'fund_scale')
    return fund_scale

def get_fund_scale(fund_code, trade_date):
    ''' 获取当前基金规模
    为了避免数据更新及时性的影响，选择前5期平均规模 '''
    fund_scale = read_scale_all()
    scale = fund_scale.loc[
        (fund_scale['基金代码']==fund_code)&
        (fund_scale['日期']<date_format(trade_date)),
        :].sort_values(by=['日期'], ignore_index=True)
    if scale.shape[0] < 5:
        return scale['期末净资产（亿元）'].mean()
    else:
        return scale['期末净资产（亿元）'].iloc[-5:].mean()

def get_fund_type(fund_code):
    ''' 获取单个基金的基金类型 '''
    info_all = read_info_all()
    info = info_all.loc[info_all['基金代码']==fund_code, '基金类型']
    if info.shape[0] > 0:
        return info.values[0]
    else:
        raise ValueError(f'查无此基：{fund_code}')

def get_fund_name(fund_code):
    ''' 获取单个基金的基金类型 '''
    info_all = read_info_all()
    info = info_all.loc[info_all['基金代码']==fund_code, '基金简称']
    if info.shape[0] > 0:
        return info.values[0]
    else:
        raise ValueError(f'查无此基：{fund_code}')

def get_fund_return(fund_code, riskfree=0):
    ''' 获取单个基金的收益率（日）序列，无风险收益率（年化）默认为0
    为了方便后续计算，缺失的净值会按前值补全
    可能会高估周净值序列的波动率 '''
    nav_all = read_nav_all()
    fund_nav = nav_all.loc[
        (nav_all['基金代码']==fund_code),
        ['日期', '累计净值']
        ].reset_index(drop=True)
    if fund_nav.shape[0] > 0:
        trade_dates = pd.DataFrame({'日期': get_trade_dates(
            fund_nav['日期'].min(),
            fund_nav['日期'].max())})
        nav_merge = pd.merge(trade_dates, fund_nav, how='left', on=['日期'])
        nav_merge['累计净值'].fillna(method='ffill', inplace=True)
        nav_merge['收益率'] = (
            nav_merge['累计净值'] / nav_merge['累计净值'].shift(1) - 1
            ) - riskfree / 250
        return nav_merge[['日期', '收益率']].iloc[1:].reset_index(drop=True)
    else:
        # 后端收费基金可能没有净值
        return pd.DataFrame(columns=['日期','收益率'])

def calc_annual_return(ret_df, start_date, end_date):
    ''' 根据收益率（日）序列，计算年化收益率 '''
    ret = ret_df.loc[
        (ret_df['日期'] >= date_format(start_date))&
        (ret_df['日期'] <= date_format(end_date)),
        '收益率'].mean()
    return ret * 250

def calc_annual_volatility(ret_df, start_date, end_date):
    ''' 根据收益率（日）序列，计算年化波动率 '''
    ret = ret_df.loc[
        (ret_df['日期'] >= date_format(start_date))&
        (ret_df['日期'] <= date_format(end_date)),
        '收益率'].std()
    return ret * np.sqrt(250)

def calc_annual_dd_vol(ret_df, start_date, end_date):
    ''' 根据收益率（日）序列，计算下行波动率 '''
    ret = ret_df.loc[
        (ret_df['日期'] >= date_format(start_date))&
        (ret_df['日期'] <= date_format(end_date)),
        '收益率'].apply(lambda x: min(x, 0)).std()
    return ret * np.sqrt(250)

def calc_drawdown(ret_df, start_date, end_date):
    ''' 计算回撤期和回撤 '''
    ret = ret_df.loc[
        (ret_df['日期'] >= date_format(start_date))&
        (ret_df['日期'] <= date_format(end_date)),
        ['日期', '收益率']].reset_index(drop=True)
    ret['虚拟净值'] = (ret['收益率'] + 1).cumprod()
    down_max = 0
    down_start = date_format(start_date)
    down_end = date_format(end_date)
    for _, row in ret.iterrows():
        nav_max = max(1, ret.iloc[:_+1, :]['虚拟净值'].max())
        down = row['虚拟净值'] / nav_max - 1
        if down <= down_max:
            down_max = down
            down_dates = ret.loc[
                ret['虚拟净值']==nav_max
                , '日期']
            if down_dates.shape[0] > 0:
                down_start = down_dates.iloc[0]
            else:
                down_dates = ret_df['日期'].iloc[0]
            down_end = row['日期']
    return (down_max, down_start, down_end)

class Fund:
    ''' 父类：基金 '''
    def __init__(
        self, fund_code, start_date=None, end_date=None,
        fund_type=None, fund_name=None, riskfree=0):
        ''' 指定基金代码、收益率起始日期、收益率终止日期、基金类型 '''
        self.fund_code = fund_code
        self.fund_ret = get_fund_return(fund_code,riskfree=riskfree)
        self.fund_type = fund_type if fund_type else get_fund_type(self.fund_code)
        self.fund_name = fund_name if fund_name else get_fund_name(self.fund_code)
        if self.fund_ret.shape[0] > 0:
            self.first_date = self.fund_ret['日期'].min()
            self.last_date = self.fund_ret['日期'].max()
            self.start_date = (
                max(self.first_date, date_format(start_date))
                if start_date else self.first_date)
            self.end_date = (
                min(self.last_date, date_format(end_date))
                if end_date else self.last_date)
            self.fund_scale = get_fund_scale(fund_code, self.end_date)
            self.annual_return = calc_annual_return(
                self.fund_ret, self.start_date, self.end_date)
            self.annual_volatility = calc_annual_volatility(
                self.fund_ret, self.start_date, self.end_date)
            self.annual_dd_vol = calc_annual_dd_vol(
                self.fund_ret, self.start_date, self.end_date)
            self.drawdown, self.drawdown_start, self.drawdown_end = calc_drawdown(
                self.fund_ret, self.start_date, self.end_date)
            self.sharpe = (
                self.annual_return / self.annual_volatility
                if self.annual_volatility > 0 else np.nan)
            self.sortino = (
                self.annual_return / self.annual_dd_vol
                if self.annual_dd_vol > 0 else np.nan)
            self.calmar = (
                self.annual_return / self.drawdown
                if self.drawdown < 0 else np.nan)
        else:
            # 应对收益率序列为空的情形
            self.first_date = pd.NaT
            self.last_date = pd.NaT
            self.start_date = pd.NaT
            self.end_date = pd.NaT
            self.fund_scale = np.nan
            self.annual_return = np.nan
            self.annual_volatility = np.nan
            self.annual_dd_vol = np.nan
            self.drawdown = np.nan
            self.drawdown_start = pd.NaT
            self.drawdown_end = pd.NaT
            self.sharpe = np.nan
            self.sortino = np.nan
            self.calmar = np.nan

if __name__ == '__main__':
    fund = Fund('377240')
    # pass